#Ana Loureiro 104063

import numpy as np
import matplotlib.pyplot as plt


def aceleracao_gravitacional(r, G, M, m):
    distancia = np.linalg.norm(r)
    return -G * M * m * r / distancia**3

def func_epg (G, M, r, m):
    distancia = np.linalg.norm(r)
    return -G * M * m / distancia

def func_ec(m,v):
    return 0.5*m*v**2


def metodo_euler_cromer_orbita(r0, v0, G, M, t0, tf, dt, m):

    n = int((tf - t0) / dt)
    t = np.linspace(t0, tf, n)
    r = np.empty((n, 2))
    v = np.empty((n, 2))
    epg = np.empty((n, 2))
    ec = np.empty((n, 2))
    em = np.empty((n, 2)) 

    r[0] = r0.copy()
    v[0] = v0.copy()

    for i in range(n - 1):
        #b)
        epg[i + 1] = func_epg(G,M,r[i],m)
        ec[i + 1] = func_ec(m,v[i])
        em[i + 1] = epg[i] + ec[i]

        a = aceleracao_gravitacional(r[i], G, M, m)
        v[i + 1] = v[i] + a * dt
        r[i + 1] = r[i] + v[i + 1] * dt



    return t, r, v , epg, ec, em


def plotA(r):
   
    plt.figure()
    plt.plot(r[:, 0], r[:, 1], label= "orbita terra")
    plt.plot(0,0, 'o', label= "posicao sol")
    plt.xlabel("x (AU)")
    plt.ylabel("y (AU)")
    plt.title("Órbitas via Método de Euler-Cromer")
    plt.grid()
    plt.legend()
    plt.axis('equal')
    plt.show()

def plotB(epg, ec, em, t):
   
    plt.figure()
    plt.plot(t, epg, label= "Epg")
    plt.plot(t, ec, label= "Ec")
    plt.plot(t, em, label= "Em (soma)")
    plt.xlabel("s")
    plt.ylabel("J")
    plt.title("Energias em funcao do tempo")
    plt.grid()
    plt.legend()
    plt.axis('equal')
    plt.show()

if __name__ == "__main__":


    G = 4 * np.pi**2  # AU³ / (M * ano²)
    M = 1.0           # massa do Sol em unidades solares
    m = 1.0
    dt = 0.01
    t0 = 0.0
    tf = 3.0        #anos

    r0 = np.array([1.0, 0.0])  # Posição inicial a 1 AU
    v0 = np.array([0, 2.3 * np.pi])  # Velocidade inicial 


##### a) a trajetoria tem forma de um circulo.
    t, r, v, epg, ec, em = metodo_euler_cromer_orbita(r0, v0, G, M, t0, tf, dt, m)
    plotA(r)


### b) A energia Cinetia foi positiva, porem sendo a energia potencia gravitica negativa numa escala muito maior, a energia total é tambem negativa
    plotB(epg, ec, em, t)